#include<iostream>
#include<map>
#include<stack>
#include<unordered_map>
#include<unordered_set>
#include<unistd.h>
#include<stdlib.h>
#include<stddef.h>
#include<string.h>
#include<sys/stat.h>
#include<errno.h>
#include"pin.H"
#include"portability.H"
#include"MEMREF.h"

#if defined(__GNUC__)
#include<cxxabi.h>
#endif

using namespace std;

//--------------------------ANALYSIS ROUTINES------------------------------

//KNOBs: eliminate redundant MEMREFs, default = ture.
KNOB<bool> KnobRedundantCheck(KNOB_MODE_WRITEONCE, "pintool", "r", "1",
	"eliminate redundant MEMREFs");
bool RedundantCheck = true;

//KNOBs: turn off/on logging, for debugging purpose.
KNOB<bool> KnobLog(KNOB_MODE_WRITEONCE, "pintool", "l", "1",
	"turn off loging");
bool LogOff = false;

//KNOBs: buffer size for every loger, default = 512 * sizoef(MEMREF).
KNOB<size_t> KnobMemrefPerLoger(KNOB_MODE_WRITEONCE, "pintool", "b", "1024",
	"every loger can buffer n MEMREFs");
size_t BufSize = 512;
const size_t BufMin = 100;
const size_t BufMax = 1024 * 1024; // sizeof(MEMREF) MB.

//KNOBs: write into file in binary format, thus saving disk space, default = false.
//we now simply write the MEMREF structure into the file, this structure can be further
//reduced, e.g. by encoding the read and size fields in the addr field since only 48 bits
//are used in a pointer on x86_64 Linux.
KNOB<bool> KnobWriteBinary(KNOB_MODE_WRITEONCE, "pintool", "B", "0",
	"write into file in  binary format");
bool WriteBinary = false;

//standard library routines are NOT thread-safe under PIN, use a lock to protect them.
PIN_LOCK _fileLock;

//every thread has a private loger to avoid lock contention. this loger is refered
//to by a pointer in the TLS structure(see below).
class Loger {
  public:
	Loger(size_t bufMax = BufSize);

	void PutText(const string &txt) {
	  _logfile << txt;
	}
	//put a MEMREF in log file.
	void PutMemRef(VOID *pc, VOID *addr, size_t size,
		int read, THREADID tid, VOID *lock);

	//put a boundary in log file.
	void PutBoundary(MEMREF::BoundaryType bType, const string &data);

	//start a PS, open a new log file for it.
	void PSStart(const string &PSDir, unsigned int threadID) {
	  assert(!_logfile.is_open());
	  char buf[20];
	  snprintf(buf, 20, "THREAD%u", threadID);
	  GetLock(&_fileLock, 1);
	  openFile(PSDir + buf);
	  ReleaseLock(&_fileLock);
	  assert(_logfile.is_open());
	}
	void PSEnd() {
	  Flush();
	  _logfile.close();
	}

	//when entering a nested PS, we need to temporaily stop current PS.
	//Stop will flush and close the log file, which will be reopened later
	//when exiting the nested PS.
	void Stop() {
	  assert(_logfile.is_open());
	  PSEnd();
	}

	//when exiting a nested PS, reopen the log file for the enclosing PS.
	void ReOpen(const string &fname) {
	  assert(!_logfile.is_open());
	  GetLock(&_fileLock, 1);
	  openFile(fname);
	  ReleaseLock(&_fileLock);
	}

	void Flush();
	~Loger() {
	  assert(!_logfile.is_open());
	  delete [] _buf;
	}

  private:
	void openFile(const string &filename) {
	  assert(!_logfile.is_open());
	  try {
		ios_base::openmode mode = fstream::out | fstream::app;
		if (WriteBinary)
		  mode |= fstream::binary;
		_logfile.open(filename.c_str(), mode);
	  } catch (...) {
		throw;
	  }
	}

	//underlying file.
	ofstream _logfile;
	//start, current postion and end of the buffer.
	MEMREF *_buf;
	MEMREF *_cursor;
	MEMREF *_end;
};

Loger::Loger(size_t bufMax) : _buf(new MEMREF[bufMax]), _cursor(_buf),
  _end(_buf + bufMax) {
  }

//put a MEMREF in the buffer, which is flushed when full.
void Loger::PutMemRef(VOID *pc, VOID *addr, size_t size,
	int read, THREADID tid, VOID *lock) {
  assert(_logfile.is_open());
  if (_cursor != _buf) {
	MEMREF &prev = *(_cursor - 1);
	if ((char *)prev.addr + prev.size == addr && prev.read == read) {
	  prev.size += size;
	  if (unlikely(prev.size > (1UL << 27))) {
		printf("too long a range-----------------------------------\n");
	  }
	  return;
	}
  }
  if (unlikely(_cursor == _end)) {
	Flush();
	assert(_cursor == _buf);
	_cursor = _buf;
  }
  _cursor->pc = pc;
  _cursor->addr = addr;
  _cursor->read = read;
  _cursor->size = size;
  _cursor->tid = tid;
  _cursor->lock = lock;
  _cursor++;
}

void Loger::PutBoundary(MEMREF::BoundaryType bType, const string &data) {
  assert(_logfile.is_open());
  Flush();
  if (WriteBinary) PutBinaryBoundary(_logfile, bType, data);
  else PutTextBoundary(_logfile, bType, data);
  _logfile.flush();

}

//demangle C++ names, for debuging.
string Demangle(const char *str) {
#if defined(__GNUC__)
  int status;
  char *demangled = __cxxabiv1::__cxa_demangle(str, 0, 0, &status);
  if (demangled) {
	string res = demangled;
	free(demangled);
	return res;
  }
#endif
  return str;
}

//get function name from PC, for debuging.
string PcToRtnName(ADDRINT pc) {
  string ret;
  PIN_LockClient();
  RTN rtn = RTN_FindByAddress(pc);
  if (RTN_Valid(rtn))
	ret = Demangle(RTN_Name(rtn).c_str());
  PIN_UnlockClient();
  return ret;
}

//flush the buffer into file.
void Loger::Flush() {
  assert(_logfile.is_open());
  char buf[500];
  if (_cursor != _buf) {
	for (MEMREF *cur = _buf; cur < _cursor; ++cur) {
	  //write to the underlying file.
	  if (WriteBinary)
		PutBinaryMEMREF(_logfile, *cur);
	  else {
		snprintf(buf, 100, "%-15p %-15p %-2u %-5s %-5u %-15p",
			cur->pc, cur->addr, (unsigned int)cur->size,
			cur->read ? "READ" : "WRITE",
			cur->tid, cur->lock);
		PutTextMEMREF(_logfile, string(buf).append(PcToRtnName((ADDRINT)cur->pc)).append("\n"));
	  }
	}
	_cursor = _buf;
  }
  _logfile.flush();
}

//every PS/team has a TeamData object to store information shared between all team members.
struct TeamData {
  TeamData(THREADID leader, const string &parentDir) : _leader(leader),
  _threadFunction(NULL),  _currentPSDir(parentDir), 
  _prevTeam(NULL), _realArg(NULL) {}

  //every team has a single leader, i.e. the one which starts a team by calling GOMP_parallel_start.
  const THREADID _leader;
  //thread function executed by every team member.
  void (*_threadFunction)(void *);
  //name of this team, it's a directory name, e.g. ./PS0_in_THREAD0
  string _currentPSDir;
  //if we are in a nested PS, this is a pointer to the previous team data, else it's set to NULL;
  TeamData *_prevTeam;
  //this stores the real argument to the thread function, see Before_GOMP_team_start for detail.
  void *_realArg;
  //map a synchronization object to its last operator.
  vector<pair<void *, string> > _sync_objs;
};

struct thread_data_t;
TLS_KEY tls_key;
thread_data_t *get_tls(TLS_KEY key, THREADID tid) {
  return static_cast<thread_data_t*>(PIN_GetThreadData(tls_key, tid));
}

struct HashKey {
  //size_t pc;   //PC.
  size_t data; // encoding: READ/WRITE(1bit) + WIDTH(3bits) + UNUSED(12bits) + REFADDR(48bits).
  bool operator ==(const HashKey &oth) const {
	//return (pc == oth.pc) && (data == oth.data);
	return (data == oth.data);
  }
};

//we need to specialize hash template class in the std namespace.
namespace std {
template<>
struct hash<HashKey> {	
  typedef size_t result_type;
  typedef HashKey argument_type;
  size_t operator()(HashKey key) const {
	//sizeof(HashKey) is 8 bytes for the time being.
	if (sizeof(HashKey) == sizeof(size_t)) {
	  //just use stock hashing for long double, sizeof(long double) is 16 bytes on x86_64 Linux.
	  hash<size_t> H;
	  return H(*(size_t *)&key);
	}
	//use string hashing for longer keys?
	hash<string> h;
	char *s = (char *)&key;
	char *e = s + sizeof(HashKey);
	return h(string(s, e));
  }
};
}

//TLS data, every thread in the application program(under which PIN runs) has one of its own.
struct thread_data_t {
  thread_data_t(THREADID tid) : _tid(tid), _isLeader(false), _team(NULL), _lock(NULL),
  _threadLoger(new Loger()) {
	_NthPS.push(0);
	_NthSBB.push(0);
  }

  ~thread_data_t() { delete _threadLoger; }
  //simply forward it to the loger.
  void PutText(const string &txt) {
	_threadLoger->PutText(txt);
  }
  void PutMemRef(VOID *pc, VOID *addr, size_t size, int read) {
	_threadLoger->PutMemRef(pc, addr, size, read, _tid, _lock);
  }
  void PutBoundary(MEMREF::BoundaryType bType, const string &data) {
	_threadLoger->PutBoundary(bType, data);
  }
  void PSStop() {
	_threadLoger->Stop();
  }
  void PSResume() {
	char buf[10];
	snprintf(buf, 10, "%u", PIN_ThreadId());
	_threadLoger->ReOpen(_team->_currentPSDir + "THREAD" + buf);
	_recorded.clear();
  }
  void PSStart() {
	_threadLoger->PSStart(_team->_currentPSDir, _tid);
	_recorded.clear();
  }
  void PSEnd() {
	_threadLoger->PSEnd();
  }

  const THREADID _tid;
  bool _isLeader;
  //pointer to the team data.
  TeamData *_team;
  //if we are in a CS, _lock is the address of the lock, else _lock == NULL.
  VOID *_lock;
  //pointer to per-thread loger.
  Loger *_threadLoger;
  stack<int> _NthPS;
  stack<int> _NthSBB;
  //pointers to the begining of stack frames, every function on the call stack
  //has a distinct pointer.
  stack<ADDRINT> _frameStart;
  //hash table to remember which MEMREFs have been recorded.
  unordered_set<HashKey> _recorded;
  private:
  TeamData &Team() { return *_team; }
  //don't allow copying.
  thread_data_t(const thread_data_t &);
  thread_data_t *operator = (const thread_data_t &);
};

//called when entering a PS, this is a threadlet boundary.
//at this point, no member thread has been created. 
VOID Before_GOMP_team_start(THREADID tid, void (*fn)(void *), void **pToArg) {
  thread_data_t *data = get_tls(tls_key, tid);
  TeamData *td = new TeamData(tid, data->_team == NULL ? "./" : data->_team->_currentPSDir);
  //save the real argument to the thread function.
  td->_realArg = *pToArg;
  td->_threadFunction = fn;
  //save current team data.
  td->_prevTeam = data->_team;
  //Note that we rewrite the second argument to GOMP_parallel_start and let it point to
  //the TeamData for this team, so when the thread functions(except the one in the leader thread) start, 
  //they acutually get a pointer to the TeamData. Then all the thread functions(again, except the one in
  //the leader thread) must extract the real argument from TeamData and rewrite its sole argument with 
  //the real argument. This is the only doable way i came up with to pass TeamData to non-leader member threads.
  *pToArg = td;

  char s[30];
  snprintf(s, 30, "PS%u_in_THREAD%d", data->_NthPS.top(), tid);
  data->_NthPS.top()++;
  data->_NthPS.push(0);
  td->_currentPSDir += s;
  //td->_currentPSDir += PcToRtnName(reinterpret_cast<ADDRINT>(td->_threadFunction));
  td->_currentPSDir += "/";
  int ret;
  //create folder for this PS.
  if ((ret = rmdir(td->_currentPSDir.c_str())) && errno != ENOENT) {
	string cmd("rm -r " + td->_currentPSDir);
	if (system(cmd.c_str())) {
	  //treat failure as fatal error.
	  printf("Can not remove directory %s\n", td->_currentPSDir.c_str());
	  exit(1);
	}
  }
  if (int ret = mkdir(td->_currentPSDir.c_str(), S_IRWXU | S_IRWXG | S_IRWXO)) {
	if (ret != EEXIST) {
	  printf("Can not create directory %s", td->_currentPSDir.c_str());
	  exit(1);
	}
  }

  //if we are already in a PS, put a boundary and stop current PS.
  if (data->_team) {
	char nthsbb[8];
	snprintf(nthsbb, 8, " %d", data->_NthSBB.top());
	data->PutBoundary(MEMREF::ENTER_PS, td->_currentPSDir + nthsbb);
	data->PSStop();
  }
  data->_NthSBB.top()++;
  data->_NthSBB.push(0);
  //switch to the new team.
  data->_team = td;
  data->_isLeader = true;

  printf("%d before team start\n", tid);
}

VOID After_GOMP_team_start(THREADID tid) {
  printf("%d after team start\n", tid);
}

//this is called at the entry point of thread functions.
VOID Before_thread_function(CONTEXT *ctxt, THREADID tid,
	VOID *threadFunAddr, void **pToArg) {
  thread_data_t *data = get_tls(tls_key, tid);
  if (!data->_isLeader) {
	//extract real argument for the thread function.
	data->_team = static_cast<TeamData *>(*pToArg);
	*pToArg = data->_team->_realArg;
  }
  assert(threadFunAddr == data->_team->_threadFunction);
  data->_NthPS.top() = 0;
  data->_NthSBB.top() = 0;
  //now start a new PS.
  data->PSStart();
  data->_frameStart.push(PIN_GetContextReg(ctxt, REG_STACK_PTR));
  printf("%d before thread function: %s\n", PIN_ThreadId(),
	  PcToRtnName((ADDRINT)threadFunAddr).c_str());
}

//this is called at the exit point of thread functions.
VOID After_thread_function(CONTEXT *ctxt, THREADID tid) {
  thread_data_t *data = get_tls(tls_key, tid);
  //end the PS.
  data->PSEnd();
  assert(data->_frameStart.top() == PIN_GetContextReg(ctxt, REG_STACK_PTR));
  data->_frameStart.pop();
  printf("%d after thread function: %s\n", tid,
	  PcToRtnName((ADDRINT)data->_team->_threadFunction).c_str());
  if (!data->_isLeader)
	data->_team = NULL;
}

//called when exiting a PS, this is a threadlet boundary.
//OpenMP garantees all threads have finished doing work at this point.
VOID After_GOMP_team_end(THREADID tid) {
  thread_data_t *data = get_tls(tls_key, tid);
  //this should be the leader thread.
  assert(data->_isLeader);
  data->_NthPS.pop();
  data->_NthSBB.pop();
  TeamData *td = data->_team;
  //switch back to previous team.
  data->_team = data->_team->_prevTeam;
  if (data->_team) {
	//resume the PS we stopped.
	data->PSResume();
	char nthsbb[8];
	snprintf(nthsbb, 8, " %d", data->_NthSBB.top());
	data->PutBoundary(MEMREF::LEAVE_PS, td->_currentPSDir + nthsbb);
	data->_NthSBB.top()++;
  }
  delete td;
  //this thread mgiht not be the leader of the previous team.
  if (data->_team && tid == data->_team->_leader)
	data->_isLeader = true;
  else
	data->_isLeader = false;
  printf("%d after team end\n", tid);
}

//called when entering a CS.
VOID Before_GOMP_critical_start(THREADID tid, VOID *lock) {
}

VOID After_GOMP_critical_start(THREADID tid, VOID *lock) {
  printf("%d after critical start\n", tid);
  thread_data_t *data = get_tls(tls_key, tid);
  assert(data->_lock == NULL);
  data->_lock = lock;
  //clear the redundant-checking set.
  data->_recorded.clear();
  TeamData *td = data->_team;
  auto it = begin(td->_sync_objs);
  for (; it != end(td->_sync_objs); it++) {
	if (it->first == data->_lock)
	  break;
  }

  //char b[20];
  //snprintf(b, 20, " %p %d", lock, data->_NthSBB.top());

  // recieve.
  if (it == td->_sync_objs.end())
	//This is the first locking operation on this lock, ignore it.
	//nothing to "receive".
	;
  else {
	data->PutBoundary(MEMREF::ENTER_CS, it->second);
	data->_NthSBB.top()++;
  }
}

static string SBBID(const string &psdir, int tid, int nsbb) {
  char buf[20];
  snprintf(buf, 20, "THREAD%d %d", tid, nsbb);
  return psdir + buf;
}
//called when exiting a CS.
VOID Before_GOMP_critical_end(THREADID tid) {
  printf("%d before critical end\n", tid);
  thread_data_t *data = get_tls(tls_key, tid);
  data->_recorded.clear();
  //char b[20];
  //snprintf(b, 20, " %p %d", data->_lock, data->_NthSBB.top());
  // post.
  data->PutBoundary(MEMREF::LEAVE_CS, "");
  TeamData *td = data->_team;
  auto it = begin(td->_sync_objs);
  for (auto it = begin(td->_sync_objs); it != end(td->_sync_objs); it++) {
	if (it->first == data->_lock)
	  break;
  }
  if (it == td->_sync_objs.end()) {
	td->_sync_objs.push_back(pair<void *, string>(data->_lock, string()));
	it = td->_sync_objs.end() - 1;
  }
  it->second = SBBID(td->_currentPSDir, data->_tid, data->_NthSBB.top());
  data->_NthSBB.top()++;
  data->_lock = NULL;
}

#define LEN 4
#define SIZEMASK ((1UL << LEN) - 1) //3 least significant bits.
void CheckAndLog(thread_data_t *tdata, VOID *pc, VOID *addr,
	size_t size, bool read) {
  if (unlikely(!RedundantCheck)) {
	tdata->PutMemRef(pc ,addr, size, read ? 1: 0);
	return;
  }
  //encode all the feilds in two words.
  HashKey key;
  //key.pc = (size_t)pc;
  assert(((size_t)addr & ~((1UL << 48) - 1)) == 0);
  key.data = (size_t)addr;
  key.data |= (size & SIZEMASK) << 59;
  if (read)
	key.data |= (1UL << 63);
  if (tdata->_recorded.count(key) == 0) {
	tdata->PutMemRef(pc ,addr, size, read ? 1: 0);
	tdata->_recorded.insert(key);
  }
}

inline void DoLog(CONTEXT *ctxt, VOID *pc, VOID *addr, size_t size, THREADID tid, bool read) {
  if (LogOff)
	return;
  thread_data_t *data = get_tls(tls_key, tid);
  if (!data->_team)   //this MEMREF is not made in any PS.
	return;
  assert(tid == data->_tid);
  ADDRINT esp = PIN_GetContextReg(ctxt, REG_STACK_PTR);
  assert(!data->_frameStart.empty());
  //check if it's stack access.
  if ((ADDRINT)addr < data->_frameStart.top() || (ADDRINT)addr > esp)
	CheckAndLog(data, pc, addr, size, read);
}

//memory reference callbacks.
VOID MemRead(CONTEXT *ctxt, VOID *pc, VOID *addr, size_t size, THREADID tid) {
  DoLog(ctxt, pc, addr, size, tid, true);
}

VOID MemWrite(CONTEXT *ctxt, VOID *pc, VOID *addr, size_t size, THREADID tid) {
  DoLog(ctxt, pc, addr, size, tid, false);
}

//this is called before every direct call instruction. This is where we caculate the start of the
//stack frame for the upcoming call, i.e. ESP - sizeof(EIP).
VOID BeforeCall(CONTEXT *ctxt, THREADID tid, ADDRINT cpc, ADDRINT pc) {
  thread_data_t *data = get_tls(tls_key, tid);
  data->_frameStart.push(PIN_GetContextReg(ctxt, REG_STACK_PTR) - sizeof(ADDRINT));
}

VOID AfterCall(CONTEXT *ctxt, THREADID tid, ADDRINT pc) {
  thread_data_t *data = get_tls(tls_key, tid);
  assert(data->_frameStart.top() == PIN_GetContextReg(ctxt, REG_STACK_PTR));
  data->_frameStart.pop();
}

//thread creation/destroying callback, called when a new thread is created/destroyed.
VOID ThreadStart(THREADID tid, CONTEXT *ctxt, INT32 flags, VOID *v) {
  thread_data_t *data = new thread_data_t(tid);
  PIN_SetThreadData(tls_key, data, tid);
}

VOID ThreadFini(THREADID tid, const CONTEXT *ctxt, INT32 flags, VOID *v) {
  thread_data_t *data = get_tls(tls_key, tid);
  assert(data->_NthPS.size() == 1);
  delete data;
  PIN_SetThreadData(tls_key, NULL, tid);
  printf("thread %d finished\n", tid);
}

//--------------------------INSTRUMENTATION ROUTINES------------------------------

//macros to save key strokes.
#define INSERTCALL(rtn, pos, which) \
  RTN_InsertCall((rtn), (pos), (AFUNPTR)(which), IARG_THREAD_ID,\
	  IARG_END)

#define INSERTCALL_FUNCARG_1(rtn, pos, which, argType, argVal) \
  RTN_InsertCall((rtn), (pos), (AFUNPTR)(which),\
	  IARG_THREAD_ID,\
	  (argType),(argVal),\
	  IARG_END)

#define INSERTCALL_FUNCARG_2(rtn, pos, which, argType1, argVal1, argType2, argVal2) \
  RTN_InsertCall((rtn), (pos), (AFUNPTR)(which),\
	  IARG_THREAD_ID,\
	  (argType1),(argVal1),\
	  (argType2),(argVal2),\
	  IARG_END)

//Functions with names like REPLACEMENT_XXX are replacement for GOMP_XXX functions.
//instrument 'parallel' construct. 
void REPLACEMENT_parallel_start(CONTEXT *ctxt, AFUNPTR Original_parrallel_start,
	void (*f) (void *), void **data, unsigned nthreads, THREADID tid) {
  Before_GOMP_team_start(tid, f, data);
  int ret = 1;
  PIN_CallApplicationFunction(ctxt, tid, CALLINGSTD_DEFAULT, Original_parrallel_start, 
	  PIN_PARG(int), &ret, 
	  PIN_PARG(void (*) (void *)), f, 
	  PIN_PARG(void *), *data, 
	  PIN_PARG(unsigned), nthreads, 
	  PIN_PARG_END());
  After_GOMP_team_start(tid);
}

void INSTRUMENT_parallel_start(RTN &rtn, IMG &img) {
  //temperally close rtn for RTN_ReplaceSignature.
  RTN_Close(rtn);
  PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_DEFAULT, "GOMP_parallel_start",
	  PIN_PARG(void (*) (void*)), PIN_PARG(void *), PIN_PARG(unsigned), PIN_PARG_END());
  RTN_ReplaceSignature(rtn, AFUNPTR(REPLACEMENT_parallel_start),
	  IARG_PROTOTYPE, proto, 
	  IARG_CONST_CONTEXT,
	  IARG_ORIG_FUNCPTR, 
	  IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
	  IARG_FUNCARG_ENTRYPOINT_REFERENCE, 1,
	  IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
	  IARG_THREAD_ID,
	  IARG_END);
  RTN_Open(rtn);
}

void REPLACEMENT_parallel_end(CONTEXT *ctxt, AFUNPTR Original_parrallel_end, THREADID tid) {
  PIN_CallApplicationFunction(ctxt, tid, CALLINGSTD_DEFAULT, Original_parrallel_end, 
	  PIN_PARG_END());
  After_GOMP_team_end(tid);
}

void INSTRUMENT_parallel_end(RTN &rtn, IMG &img) {
  RTN_Close(rtn);
  RTN_ReplaceSignature(rtn, AFUNPTR(REPLACEMENT_parallel_end), IARG_CONST_CONTEXT,
	  IARG_ORIG_FUNCPTR, 
	  IARG_THREAD_ID,
	  IARG_END);
  RTN_Open(rtn);
}

//instrument 'parallel for' construct.
#define REPLACEMENT_PARALLEL_LOOP_START(schedule) \
  void REPLACEMENT_parallel_loop_##schedule##_start(CONTEXT *ctxt,\
	  AFUNPTR Original_parallel_loop_##schedule##_start,\
	  void (*f) (void *), void **data, unsigned nthreads,\
	  long start, long end, long incr, long chunk_size,\
	  THREADID tid) {\
	Before_GOMP_team_start(tid, f, data);\
	int ret = 1;\
	PIN_CallApplicationFunction(ctxt, tid, CALLINGSTD_DEFAULT,\
		Original_parallel_loop_##schedule##_start,\
		PIN_PARG(int), &ret,\
		PIN_PARG(void (*) (void *)), f,\
		PIN_PARG(void *), *data,\
		PIN_PARG(unsigned), nthreads,\
		PIN_PARG(long), start, PIN_PARG(long), end, PIN_PARG(long), incr,\
		PIN_PARG(long), chunk_size,\
		PIN_PARG_END());\
	After_GOMP_team_start(tid);\
  }

#define INSTRUMENT_PARALLEL_LOOP_START(schedule)\
  void INSTRUMENT_parallel_loop_##schedule##_start(RTN &rtn, IMG &img) {\
	RTN_Close(rtn);\
	PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_DEFAULT,\
		"GOMP_parallel_loop_"#schedule"_start",\
		PIN_PARG(void (*) (void*)), PIN_PARG(void *), PIN_PARG(unsigned),\
		PIN_PARG(long), PIN_PARG(long), PIN_PARG(long), PIN_PARG(long),\
		PIN_PARG_END());\
	RTN_ReplaceSignature(rtn, AFUNPTR(REPLACEMENT_parallel_loop_##schedule##_start),\
		IARG_PROTOTYPE, proto,\
		IARG_CONST_CONTEXT,\
		IARG_ORIG_FUNCPTR,\
		IARG_FUNCARG_ENTRYPOINT_VALUE, 0,\
		IARG_FUNCARG_ENTRYPOINT_REFERENCE, 1,\
		IARG_FUNCARG_ENTRYPOINT_VALUE, 2,\
		IARG_FUNCARG_ENTRYPOINT_VALUE, 3,\
		IARG_FUNCARG_ENTRYPOINT_VALUE, 4,\
		IARG_FUNCARG_ENTRYPOINT_VALUE, 5,\
		IARG_FUNCARG_ENTRYPOINT_VALUE, 6,\
		IARG_THREAD_ID,\
		IARG_END);\
	RTN_Open(rtn);\
  }

REPLACEMENT_PARALLEL_LOOP_START(static);
REPLACEMENT_PARALLEL_LOOP_START(dynamic);
REPLACEMENT_PARALLEL_LOOP_START(guided);
INSTRUMENT_PARALLEL_LOOP_START(static);
INSTRUMENT_PARALLEL_LOOP_START(dynamic);
INSTRUMENT_PARALLEL_LOOP_START(guided);

void REPLACEMENT_parallel_loop_runtime_start(CONTEXT *ctxt,
	AFUNPTR Original_parallel_loop_runtime_start,
	void (*f) (void *), void **data, unsigned nthreads,
	long start, long end, long incr,
	THREADID tid) {
  Before_GOMP_team_start(tid, f, data);
  int ret = 1;
  PIN_CallApplicationFunction(ctxt, tid, CALLINGSTD_DEFAULT, Original_parallel_loop_runtime_start,
	  PIN_PARG(int), &ret,
	  PIN_PARG(void (*) (void *)), f,
	  PIN_PARG(void *), *data,
	  PIN_PARG(unsigned), nthreads,
	  PIN_PARG(long), start, PIN_PARG(long), end, PIN_PARG(long), incr,
	  PIN_PARG_END());
  After_GOMP_team_start(tid);
}

void INSTRUMENT_parallel_loop_runtime_start(RTN &rtn, IMG &img) {
  RTN_Close(rtn);
  PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_DEFAULT,
	  "GOMP_parallel_loop_runtime_start",
	  PIN_PARG(void (*) (void*)), PIN_PARG(void *), PIN_PARG(unsigned),
	  PIN_PARG(long), PIN_PARG(long), PIN_PARG(long),
	  PIN_PARG_END());
  RTN_ReplaceSignature(rtn, AFUNPTR(REPLACEMENT_parallel_loop_runtime_start),
	  IARG_PROTOTYPE, proto,
	  IARG_CONST_CONTEXT,
	  IARG_ORIG_FUNCPTR,
	  IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
	  IARG_FUNCARG_ENTRYPOINT_REFERENCE, 1,
	  IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
	  IARG_FUNCARG_ENTRYPOINT_VALUE, 3,
	  IARG_FUNCARG_ENTRYPOINT_VALUE, 4,
	  IARG_FUNCARG_ENTRYPOINT_VALUE, 5,
	  IARG_THREAD_ID,
	  IARG_END);
  RTN_Open(rtn);
}

//instrument 'critical' construct.
void REPLACEMENT_critical_start(CONTEXT *ctxt, AFUNPTR Original_critical_start,
	void *lock, THREADID tid) {
  Before_GOMP_critical_start(tid, lock);
  PIN_CallApplicationFunction(ctxt, tid, CALLINGSTD_DEFAULT, Original_critical_start,
	  PIN_PARG_END());
  After_GOMP_critical_start(tid, lock);
}

void INSTRUMENT_critical_start(RTN &rtn, IMG &img) {
  RTN_Close(rtn);
  RTN_ReplaceSignature(rtn, AFUNPTR(REPLACEMENT_critical_start),
	  IARG_CONST_CONTEXT,
	  IARG_ORIG_FUNCPTR, 
	  IARG_PTR, 1,
	  IARG_THREAD_ID,
	  IARG_END);
  RTN_Open(rtn);
}

void INSTRUMENT_critical_end(RTN &rtn, IMG &img) {
  INSERTCALL(rtn, IPOINT_BEFORE, Before_GOMP_critical_end);
}

//instrument OpenMP thread function generated by compilers.
typedef unordered_map<string, void (*)(RTN &rtn, IMG &img)> HASH_MAP;
HASH_MAP Rtn2Instrumentor;

typedef unordered_set<string> HASH_SET;
HASH_SET Instrumented;

//depth-fisrt search the call graph and instrument every function in the *main* image,
//those in the dynamic libraries are ignored.
void WalkAndInstrumentCallGraph(RTN &rtn, IMG &img) {
  string name = RTN_Name(rtn);
  //we've visited this one.
  if (Instrumented.count(name))
	return;
  printf("instrument %s\n", name.c_str());
  Instrumented.insert(name);
  if (name.find("omp_fn.") != string::npos) {
	//this is a thread function.
	RTN_InsertCall(rtn, IPOINT_BEFORE, (AFUNPTR)Before_thread_function,
		IARG_CONST_CONTEXT,
		IARG_THREAD_ID,
		IARG_ADDRINT, RTN_Address(rtn),
		IARG_FUNCARG_ENTRYPOINT_REFERENCE, 0,
		IARG_END);
	RTN_InsertCall(rtn, IPOINT_AFTER, (AFUNPTR)After_thread_function,
		IARG_CONST_CONTEXT,
		IARG_THREAD_ID,
		IARG_END);
  }
  //instrument every memOP and collect direct calls.
  vector<ADDRINT> calledFun;
  for (INS ins = RTN_InsHead(rtn); INS_Valid(ins); ins = INS_Next(ins)) {
	//only handle direct calls.
	if (INS_IsDirectCall(ins)) {
	  ADDRINT target = INS_DirectBranchOrCallTargetAddress(ins);
	  string targetFunName = RTN_FindNameByAddress(target);
	  RTN targetRtn = RTN_FindByName(img, targetFunName.c_str());
	  if (!RTN_Valid(targetRtn)) //highly likely it's from some dynamic libraries.
		printf("skip %x %s\n", (int)INS_Address(ins), targetFunName.c_str());
	  if (RTN_Valid(targetRtn) && Instrumented.count(targetFunName) == 0) {
		if (targetFunName.find("omp_fn.") == string::npos) {
		  //some parameters are passed for debuging purpose.
		  //TODO::remove unnecessary paremeters.
		  INS_InsertPredicatedCall(ins, IPOINT_BEFORE, (AFUNPTR)BeforeCall,
			  IARG_CONST_CONTEXT,
			  IARG_THREAD_ID,
			  IARG_ADDRINT, INS_Address(ins),
			  IARG_ADDRINT, target,
			  IARG_END);
		  INS_InsertPredicatedCall(ins, IPOINT_TAKEN_BRANCH, (AFUNPTR)AfterCall,
			  IARG_CONST_CONTEXT,
			  IARG_THREAD_ID,
			  IARG_ADDRINT, INS_Address(ins),
			  IARG_END);
		  calledFun.push_back(target);
		}
	  }
	  continue;
	}
	//skip stack access, non-direct calls, returns and system calls.
	if (INS_IsStackRead(ins) || INS_IsStackWrite(ins) ||
		INS_IsCall(ins) || INS_IsRet(ins) || INS_IsSyscall(ins))
	  continue;
	UINT32 memoryOperands = INS_MemoryOperandCount(ins);
	if (memoryOperands) {
	  for (UINT32 memOp = 0; memOp < memoryOperands; memOp++) {
		UINT32 refSize = INS_MemoryOperandSize(ins, memOp);
		//note that if the operand is both read and written we log it once for each.
		if (INS_MemoryOperandIsRead(ins, memOp)) {
		  INS_InsertPredicatedCall(ins, IPOINT_BEFORE, (AFUNPTR)MemRead,
			  IARG_CONST_CONTEXT,
			  IARG_INST_PTR,
			  IARG_MEMORYOP_EA, memOp,
			  IARG_UINT32, refSize, 
			  IARG_THREAD_ID,
			  IARG_END);
		}
		if (INS_MemoryOperandIsWritten(ins, memOp)) {
		  INS_InsertPredicatedCall(ins, IPOINT_BEFORE, (AFUNPTR)MemWrite,
			  IARG_CONST_CONTEXT,
			  IARG_INST_PTR,
			  IARG_MEMORYOP_EA, memOp,
			  IARG_UINT32, refSize, 
			  IARG_THREAD_ID,
			  IARG_END);
		}
	  }
	}
  }
  RTN_Close(rtn);
  //now recursively visit functions called by this function.
  for (auto it = begin(calledFun); it != end(calledFun); ++it) {
	RTN called = RTN_FindByAddress(*it);
	assert(RTN_Valid(called));
	RTN_Open(called);
	WalkAndInstrumentCallGraph(called, img);
	RTN_Close(called);
  }
  RTN_Open(rtn);
  printf("instrument %s done\n", name.c_str());
}

void BuildInstrumentationMap() {
  //'parallel' construct.
  Rtn2Instrumentor["GOMP_parallel_start"] = &INSTRUMENT_parallel_start;
  Rtn2Instrumentor["GOMP_parallel_end"] = &INSTRUMENT_parallel_end;

  //'parallel for' construct.
  Rtn2Instrumentor["GOMP_parallel_loop_static_start"] = &INSTRUMENT_parallel_loop_static_start;
  Rtn2Instrumentor["GOMP_parallel_loop_dynamic_start"] = &INSTRUMENT_parallel_loop_dynamic_start;
  Rtn2Instrumentor["GOMP_parallel_loop_guided_start"] = &INSTRUMENT_parallel_loop_guided_start;
  Rtn2Instrumentor["GOMP_parallel_loop_runtime_start"] = &INSTRUMENT_parallel_loop_runtime_start;

  Rtn2Instrumentor["GOMP_critical_start"] = &INSTRUMENT_critical_start;
  Rtn2Instrumentor["GOMP_critical_end"] = &INSTRUMENT_critical_end;
  Rtn2Instrumentor["GOMP_thread_function"] = &WalkAndInstrumentCallGraph;
}

//instrument an image. note that we ignore all functions in the dynamic libraries for simplicity,
//which means we will miss access in those functions.
VOID Image(IMG img, VOID *v) {
  //skip dynamic libraries except libgomp.so, where GOMP_XXX functions reside.
  if (IMG_Type(img) > 2) {
	if (string::npos == IMG_Name(img).rfind("libgomp.so"))
	  return;
  }
  for (SEC sec = IMG_SecHead(img); SEC_Valid(sec); sec = SEC_Next(sec)) {
	for (RTN rtn = SEC_RtnHead(sec); RTN_Valid(rtn); rtn = RTN_Next(rtn)) {
	  RTN_Open(rtn);
	  string name = RTN_Name(rtn);
	  if (string::npos != name.find(string("omp_fn."))) {
		(Rtn2Instrumentor["GOMP_thread_function"])(rtn, img);
	  } else {
		HASH_MAP::iterator it = Rtn2Instrumentor.find(name);
		if (it != Rtn2Instrumentor.end()) {
		  (it->second)(rtn, img);
		}
	  }
	  RTN_Close(rtn);
	}
  }
}

VOID Fini(INT32 code, VOID *v) {
  printf("program finished \n");
}

int main(int argc, char **argv)
{
  PIN_InitSymbols();
  if (PIN_Init(argc, argv))
	return 1;
  InitLock(&_fileLock);
  RedundantCheck = KnobRedundantCheck.Value();
  WriteBinary = KnobWriteBinary.Value();
  LogOff = !KnobLog.Value();
  BufSize = KnobMemrefPerLoger.Value();
  if (BufSize < BufMin)
	BufSize = BufMin;
  else if (BufSize > BufMax)
	BufSize = BufMax;

  tls_key = PIN_CreateThreadDataKey(0);
  BuildInstrumentationMap();

  //add callbacks.
  PIN_AddThreadStartFunction(ThreadStart, 0);
  PIN_AddThreadFiniFunction(ThreadFini, 0);
  PIN_AddFiniFunction(Fini, 0);
  IMG_AddInstrumentFunction(Image, 0);

  PIN_StartProgram();
  return 0;
}
